import os
import re
from pathlib import Path
from dotenv import load_dotenv
from text2sql import SQLRunner, NL2SQLAgent
load_dotenv()
ROOT = Path(__file__).resolve().parents[2]

if __name__ == "__main__":
    db_url = os.getenv("DB_URL", "mysql+pymysql://用户名:密码@localhost:3306/mydatabase")
    csv_path = ROOT / "data" / "sales.csv"
    # 1) 初始化数据库与示例表
    runner = SQLRunner(db_url)
    runner.init_from_csv(str(csv_path), table_name="sales")
    # 2) 提供 schema_hint（强烈建议绑定到你企业真实数据字典）
    schema = """CREATE TABLE sales (product_name TEXT,year INTEGER,sales_amount INTEGER,region TEXT);"""
    agent = NL2SQLAgent(schema_hint=schema)

    # 3) 自然语言问题
    q = "查找2023年所有销售额超过1000万的产品，并按销售额降序排列"
    sql = agent.generate_sql(q)

    # 使用 split() 方法分割字符串，以 "</think>" 为分隔符，并取最后一部分
    parts = sql.split("</think>")
    # 提取最后一部分并去除首尾空格
    real_sql = parts[-1].strip()
    print("生成的SQL：\n", real_sql.strip())

    # 4) 执行与结果
    df = runner.exec(real_sql.strip())
    print("\n查询结果：\n", df.to_string(index=False).strip())